package com.yanti.util;

import com.yanti.domain.entity.mongodb.ChatMessage;
import com.yanti.domain.entity.mongodb.Conversation;
import com.yanti.filter.LoginUserContextHolder;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Component;

import java.time.Duration;
import java.time.LocalDateTime;
import java.util.*;
import java.util.stream.Collectors;

@Slf4j
@Component
public class MongoChatMemory implements ChatMemory {
    // 注入MongoDB 的MongoTemplate
    @Resource
    private final MongoTemplate mongoTemplate;


    public MongoChatMemory(MongoTemplate mongoTemplate) {
        this.mongoTemplate = mongoTemplate;
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        if (messages == null || messages.isEmpty()) {
            return;
        }
        conversationId = conversationId + Objects.requireNonNull(LoginUserContextHolder.getUserId());
        Query query = new Query(Criteria.where("conversationId").is(conversationId).and("deleted").is(0));
        Conversation conversation = mongoTemplate.findOne(query, Conversation.class);

        List<ChatMessage> newMessages = messages.stream()
                .map(m -> {
                    ChatMessage cm = new ChatMessage(m);
                    cm.setCreatedAt(LocalDateTime.now());
                    return cm;
                }).toList();

        if (conversation != null && conversation.getMessages() != null && !conversation.getMessages().isEmpty()) {
            ChatMessage lastMessage = conversation.getMessages().get(conversation.getMessages().size() - 1);
            // 只做和最后一条比对，避免误伤
            ChatMessage firstNew = newMessages.get(0);
            if (firstNew.getText().equals(lastMessage.getText())
                    && firstNew.getMessageType() == lastMessage.getMessageType()
                    && Duration.between(lastMessage.getCreatedAt(), firstNew.getCreatedAt()).getSeconds() < 5) {
                // 短时间内的连续重复，视为误触发
                log.warn("Skipping duplicated message: {}", firstNew.getText());
                return;
            }
            // 添加到数据库
            Update update = new Update();
            update.push("messages").each(newMessages.toArray());
            update.set("lastUpdated", new Date());
            mongoTemplate.updateFirst(query, update, Conversation.class);
        } else {
            // 新会话
            Conversation newConversation = Conversation.builder()
                    .conversationId(conversationId)
                    .assistantType("spring-ai")
                    .uid(LoginUserContextHolder.getUserId().toString())
                    .createdAt(new Date())
                    .lastUpdated(new Date())
                    .deleted(0)
                    .messages(new ArrayList<>(newMessages))
                    .build();

            mongoTemplate.insert(newConversation);
        }
    }


    @Override
    public List<Message> get(String conversationId, int lastN) {
        // 构造会话ID
        conversationId  = conversationId+Objects.requireNonNull(LoginUserContextHolder.getUserId()).toString();
        Query query = new Query(Criteria.where("conversationId").is(conversationId).and("deleted").is(0));
        Conversation conversation = mongoTemplate.findOne(query, Conversation.class);
        if (conversation == null || conversation.getMessages() == null) {
            return List.of();
        }

        List<ChatMessage> chatMessages = conversation.getMessages();
        int fromIndex = Math.max(0, chatMessages.size() - lastN);
        return chatMessages.subList(fromIndex, chatMessages.size())
                .stream()
                .map(ChatMessage::toMessage)
                .toList();
    }

    @Override
    public void clear(String conversationId) {
        // 构造会话ID
        conversationId  = conversationId+Objects.requireNonNull(LoginUserContextHolder.getUserId()).toString();
        Query query = new Query(Criteria.where("conversationId").is(conversationId));
        Update update = new Update()
                .set("messages", new ArrayList<>())
                .set("lastUpdated", new Date());
        mongoTemplate.updateFirst(query, update, Conversation.class);
    }
}
